import UIKit

class WaterfallViewController: UIViewController, UICollectionViewDataSource {
    
    // 数据源：随机生成高度
    let data = Array(1...50).map { _ in CGFloat.random(in: 100...200) }
    
    // 创建 CollectionView
    var collectionView: UICollectionView!
    
    override func viewDidLoad() {
        super.viewDidLoad()
        view.backgroundColor = .white // 背景色为白色
        
        // 初始化自定义的布局
        let layout = WaterfallLayout()
        layout.numberOfColumns = 2 // 设置为两列
        layout.cellPadding = 10 // 设置单元格的间距
        layout.delegate = self // 设置设置布局代理，用于动态计算单元格高度
        
        // 初始化 CollectionView
        collectionView = UICollectionView(frame: .zero, collectionViewLayout: layout)
        // 关闭自动布局
        collectionView.translatesAutoresizingMaskIntoConstraints = false
        collectionView.backgroundColor = .white
        // 设置数据源代理
        collectionView.dataSource = self
        // // 注册自定义单元格
        collectionView.register(WaterfallCell.self, forCellWithReuseIdentifier: "WaterfallCell")
        
        view.addSubview(collectionView)
        
        // 设置 CollectionView 布局约束
        NSLayoutConstraint.activate([
            collectionView.topAnchor.constraint(equalTo: view.safeAreaLayoutGuide.topAnchor),
            collectionView.leadingAnchor.constraint(equalTo: view.leadingAnchor),
            collectionView.trailingAnchor.constraint(equalTo: view.trailingAnchor),
            collectionView.bottomAnchor.constraint(equalTo: view.safeAreaLayoutGuide.bottomAnchor)
        ])
    }
    
    // MARK: - UICollectionViewDataSource
    func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
        return data.count // 数据的总数量
    }
    
    func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
        // 获取自定义的单元格
        let cell = collectionView.dequeueReusableCell(withReuseIdentifier: "WaterfallCell", for: indexPath) as! WaterfallCell
        cell.configure(height: data[indexPath.item]) // 配置单元格高度
        let label = UILabel()
        label.frame = cell.bounds
        label.text = String(indexPath.item)
        label.textColor = .white
        label.textAlignment = .center
        label.font = UIFont.systemFont(ofSize: 20, weight: .bold)
        cell.contentView.addSubview(label)
        return cell
    }
    // 点击处理
    func collectionView(_ collectionView: UICollectionView, didSelectItemAt indexPath: IndexPath) {
        print("第\(indexPath.row)个item被点击")
    }
}

// MARK: - 自定义 Layout
protocol WaterfallLayoutDelegate: AnyObject {
    func collectionView(_ collectionView: UICollectionView, heightForItemAt indexPath: IndexPath) -> CGFloat
}

class WaterfallLayout: UICollectionViewLayout {
    weak var delegate: WaterfallLayoutDelegate? // 代理用于获取高度
    
    var numberOfColumns = 2 // 列数
    var cellPadding: CGFloat = 10 // 单元格间距
    
    private var cache: [UICollectionViewLayoutAttributes] = [] // 缓存布局属性
    private var contentHeight: CGFloat = 0 // 内容的总高度
    
    // 获取 CollectionView 的宽度
    private var width: CGFloat {
        return collectionView?.bounds.width ?? 0
    }
    
    override func prepare() {
        guard cache.isEmpty, let collectionView = collectionView else { return }
        
        // 每列的宽度
        let columnWidth = (width - CGFloat(numberOfColumns - 1) * cellPadding) / CGFloat(numberOfColumns)
        var xOffset: [CGFloat] = [] // 每列的 X 偏移量
        for column in 0..<numberOfColumns {
            xOffset.append(CGFloat(column) * (columnWidth + cellPadding))
        }
        
        // 每列的高度
        var columnHeights = [CGFloat](repeating: 0, count: numberOfColumns)
        var column = 0 // 当前列索引
        
        // 遍历 CollectionView 的每个单元格
        for item in 0..<collectionView.numberOfItems(inSection: 0) {
            // 获取当前单元格的索引路径
            let indexPath = IndexPath(item: item, section: 0)
            // 获取单元格的高度（通过代理获取）
            let height = delegate?.collectionView(collectionView, heightForItemAt: indexPath) ?? 0
            let frame = CGRect(x: xOffset[column], y: columnHeights[column], width: columnWidth, height: height)
            let insetFrame = frame.insetBy(dx: cellPadding / 2, dy: cellPadding / 2)
            
            // 创建布局属性并缓存
            let attributes = UICollectionViewLayoutAttributes(forCellWith: indexPath)
            attributes.frame = insetFrame
            cache.append(attributes)
            
            // 更新当前列的高度
            columnHeights[column] += height + cellPadding
            
            // 找出高度最小的列作为下一个单元格的目标列
            column = columnHeights.firstIndex(of: columnHeights.min()!)!
        }
        // 更新内容总高度
        contentHeight = columnHeights.max() ?? 0
    }
    
    override var collectionViewContentSize: CGSize {
        return CGSize(width: width, height: contentHeight)
    }
    
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        return cache.filter { $0.frame.intersects(rect) }
    }
    
    override func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        return cache[indexPath.item]
    }
}

// MARK: - WaterfallLayoutDelegate 实现
extension WaterfallViewController: WaterfallLayoutDelegate {
    func collectionView(_ collectionView: UICollectionView, heightForItemAt indexPath: IndexPath) -> CGFloat {
        return data[indexPath.item] // 返回单元格的高度
    }
}

// MARK: - 自定义单元格
class WaterfallCell: UICollectionViewCell {
    override init(frame: CGRect) {
        super.init(frame: frame)
        backgroundColor = UIColor(
            red: CGFloat.random(in: 0...1),
            green: CGFloat.random(in: 0...1),
            blue: CGFloat.random(in: 0...1),
            alpha: 1.0
        ) // 设置随机背景色
        layer.cornerRadius = 8 // 设置圆角
        clipsToBounds = true
    }
    
    required init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
    
    // 配置方法（如果需要在单元格中显示内容可以扩展）
    func configure(height: CGFloat) {
        // 此处可以扩展配置逻辑
    }
}
